use num_traits::{Num, pow};
use crate::autograd::function::{ForwardArgs, Function};
use crate::autograd::function_ctx::FunctionCtx;
use crate::Tensor;

pub struct PowBackward<T> {
    ctx: FunctionCtx<T>,
}

impl<T> Function<T> for PowBackward<T>
    where
        T: Copy + Num + 'static,
{
    fn new(ctx: FunctionCtx<T>) -> Self {
        Self { ctx }
    }

    fn ctx(&self) -> &FunctionCtx<T> {
        &self.ctx
    }

    fn forward(ctx: &mut FunctionCtx<T>, args: ForwardArgs<T>) -> Tensor<T>
        where
            T: Copy + Num + 'static,
            Self: Sized + 'static,
    {
        if let ForwardArgs::TensorOther(tensor, exp) = args {
            ctx.save_tensors([tensor.clone()].into());
            ctx.save_others([exp].into());
            let exp = exp[0];
            tensor.map(|x| pow(x, exp))
        } else {
            unreachable!()
        }
    }

    fn backward(&self, grad_output: Tensor<T>) -> Vec<Option<Tensor<T>>> {
        let tensor = &self.ctx.tensors()[0];
        let exp = self.ctx.others()[0][0];
        let grad = tensor.zip(&grad_output, |a, b|
            pow(a, exp) * b);
        [grad.into()].into()
    }
}
